package com.echo.aichat.ai.agent;

import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonPropertyDescription;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.context.annotation.Description;
import org.springframework.stereotype.Component;

import java.io.File;
import java.lang.management.ManagementFactory;
import java.util.function.Function;

@Agent
@Description("评估电脑性能的智能体，输出电脑各项指标信息")
public class PerformanceEvaluator extends AbstractAgent<PerformanceEvaluator.Request, String> {
    private final String SYSTEM = """
            你是一个专业的电脑性能评估专家，可以输出电脑各项指标信息，
            如CPU使用率、内存使用情况、磁盘使用率等。
            """;

    public PerformanceEvaluator(ChatModel chatModel) {
        super(chatModel);
    }

    @Override
    public String apply(Request request) {
        return getChatClient()
                .prompt()
                .system(SYSTEM)
                .user(request.query)
                .call()
                .content();
    }

    public record Request(
            @JsonProperty(required = true)
            @JsonPropertyDescription("用户原始的提问") String query) {
    }

    @Component
    @Description("获取当前CPU使用率")
    public static class CpuUsage implements Function<CpuUsage.Request, String> {
        @Override
        public String apply(Request request) {
            com.sun.management.OperatingSystemMXBean osBean =
                    (com.sun.management.OperatingSystemMXBean) ManagementFactory.getOperatingSystemMXBean();
            double cpuLoad = osBean.getSystemCpuLoad() * 100;
            return String.format("CPU使用率: %.2f%%", cpuLoad);
        }

        public record Request(
                @JsonProperty("dummy") String dummy
        ) {}
    }

    @Component
    @Description("获取当前内存使用情况")
    public static class MemoryUsage implements Function<MemoryUsage.Request, String> {
        @Override
        public String apply(Request request) {
            com.sun.management.OperatingSystemMXBean osBean =
                    (com.sun.management.OperatingSystemMXBean) ManagementFactory.getOperatingSystemMXBean();
            long totalMemory = osBean.getTotalPhysicalMemorySize();
            long freeMemory = osBean.getFreePhysicalMemorySize();
            long usedMemory = totalMemory - freeMemory;
            return String.format("内存使用情况: 总内存: %d MB, 已用: %d MB, 空闲: %d MB",
                    totalMemory / (1024 * 1024),
                    usedMemory / (1024 * 1024),
                    freeMemory / (1024 * 1024));
        }

        public record Request(
                @JsonProperty("dummy") String dummy
        ) {}
    }

    @Component
    @Description("获取当前磁盘使用情况")
    public static class DiskUsage implements Function<DiskUsage.Request, String> {
        @Override
        public String apply(Request request) {
            // 以用户主目录所在分区作为示例
            File root = new File(System.getProperty("user.home"));
            long totalSpace = root.getTotalSpace();
            long freeSpace = root.getFreeSpace();
            long usedSpace = totalSpace - freeSpace;
            double usagePercent = (double) usedSpace / totalSpace * 100;
            return String.format("磁盘使用情况: 总空间: %d GB, 已用: %d GB (%.2f%%), 空闲: %d GB",
                    totalSpace / (1024L * 1024 * 1024),
                    usedSpace / (1024L * 1024 * 1024),
                    usagePercent,
                    freeSpace / (1024L * 1024 * 1024));
        }

        public record Request(
                @JsonProperty("dummy") String dummy
        ) {}
    }
}
